import tensorflow as tf
import numpy as np
import pickle
from self_factory import DictionaryMap
w2i, i2w = DictionaryMap.dictionary_from_pkl_file('maps.pkl')


def inner(str):
    return input(str+':')


def get_input(dictionary):
    s = inner('input')

    print('A: %s'%s)
    def str2array_func(s):
        sentences_2id = np.asarray([dictionary[word] for word in ' '.join(s).split()],dtype=np.int64)
        sentences_2id = sentences_2id.reshape([1, len(sentences_2id)])
        return sentences_2id, len(sentences_2id)
    return str2array_func(s)


def get_result(w2i, i2w):
    # 加载字典
    #dictionary, dictionary_i2w = dict_from_pkl_file('data/maps.pkl')
    dictionary = w2i
    dictionary_i2w = i2w
    # 一：导入meta文件，该文件保存图的所有信息
    saver = tf.train.import_meta_graph(r'ckpt/m.meta')

    # 二: 进入默认图
    graph = tf.get_default_graph()

    # 三：读取需要使用的tensor
    tensor_target = graph.get_tensor_by_name('placeholder/target:0')
    tensor_input = graph.get_tensor_by_name('placeholder/input:0')
    tensor_input_length = graph.get_tensor_by_name('placeholder/input_length:0')
    tensor_target_length = graph.get_tensor_by_name('placeholder/target_length:0')
    #tensor_prediction = graph.get_tensor_by_name('net_seq2seq/decoder/ArgMax:0')

    # 加载参数，必须在session中才能导入参数
    with tf.Session() as sess:
        # 四、 在session里加载参数
        # （注意：latest_checkpoint需要输入为模型保存的文件夹，因为需要meta和index文件）
        saver.restore(sess, tf.train.latest_checkpoint('ckpt/'))
        #print('get_operations:',graph.get_operations())
        #print('get_all_collection_keys', graph.get_all_collection_keys())
        print('get_name_scope',graph.get_name_scope())
        print('get_tensor_name:',[a.name for a in graph.as_graph_def().node])
        # 五 feed 给需要数据的tensor

        #data_input, len_input = get_input(dictionary)
        #feed = {tensor_input: data_input,tensor_input_length:[len_input], tensor_target_length: [len_input+10]}
        # 六 sess.run 需要拿到的结果
        #result = sess.run(tensor_prediction, feed_dict=feed)
        #print('B:', ''.join([dictionary_i2w[w[0]] for w in result]))
    # 无论是feeddict还是执行，都直接执行tensor，而不是operation
get_result(w2i, i2w)